from convdiffsolverbase import ConvDiffSolverBase
from dolfin import *
import numpy

class doCarmo(ConvDiffSolverBase):
    '''Improved SUPG method; Consistent Approximate Upwind do Carmo(2004) and do Carmo(1987).'''
    def __init__(self):
        self.name = 'doCarmo'
        ConvDiffSolverBase.__init__(self)

    def solve(self,problem,get_tau):
        '''Values obtained by get_tau are also a base for sigma_. Treatmenet of nonlinearity is based on doCarmo(2004)--no outer products.'''
        problemName, mesh, eps, BCValues, b, f, BCIndicators = problem.get_vars()

        V = FunctionSpace(mesh,'CG',1) 
        VV = VectorFunctionSpace(mesh,'CG',1)
        DGV = FunctionSpace(mesh,'DG',0)
        DGVV = VectorFunctionSpace(mesh,'DG',0)

        u = TrialFunction(V)
        varphi = TestFunction(V)
        u_ = Function(V)
      
        bcs = [DirichletBC(V,value,where) for value, where in zip(BCValues, BCIndicators)]

        v = Function(VV)
        v = interpolate(b,VV)                 # velocity

        #-----------------------------------------------------------------------------------------------------------------------------
        def projection(f,v,u,p_):
            p = TrialFunction(DGVV)
            phi = TestFunction(DGVV)

            a = inner(p,phi)*dx
            w = conditional(
                        gt(sqrt(inner(nabla_grad(u),nabla_grad(u))),1E-16),\
                        (-eps*nabla_div(nabla_grad(u))+inner(v,nabla_grad(u))-f)*nabla_grad(u)/inner(nabla_grad(u),nabla_grad(u)),\
                        Constant((0.,0.))\
                            ) 
            l = inner(w,phi)*dx
        
            solve(a==l,p_)
            
        def get_sigma(f,v,u,p_,eps,sigma_):
            p_norm = TrialFunction(DGV)
            v_norm = TrialFunction(V)
            res = TrialFunction(DGV)
            conv = TrialFunction(DGV)

            norm_v = Function(V)
            norm_p = Function(DGV)
            convF = Function(DGV)
            resF = Function(DGV)
            
            psi = TestFunction(DGV)
          
            p_norm_a =  p_norm*psi*dx
            p_norm_l = sqrt(inner(p_,p_))*psi*dx
            solve(p_norm_a==p_norm_l,norm_p)       # get the norm of p_(discon)

            v_norm_a = v_norm*varphi*dx
            v_norm_b = sqrt(inner(v,v))*varphi*dx
            solve(v_norm_a==v_norm_b,norm_v)       # get the norm of v (continuous)  
            
            resA = res*psi*dx
            resL = (-eps*nabla_div(nabla_grad(u))+inner(v,nabla_grad(u))-f)*psi*dx
            solve(resA==resL,resF)                 # get the residuum

            convA = conv*psi*dx
            convL = inner(v,nabla_grad(u))*psi*dx  # get the convective term
            solve(convA==convL,convF)
        
            for i in xrange(len(sigma_.vector().array())):
                c = Cell(mesh,i)
                x = c.midpoint()

                v_mag = norm_v(x.x(),x.y())
                p_mag = norm_p(x.x(),x.y())
                residuum = resF(x.x(),x.y())
                convective = convF(x.x(),x.y())
                
                if abs(residuum) < 1E-16:
                    sigma_.vector()[i] = 0
                else:
                    alfa = max(1,convective/residuum)
                    if p_mag >= v_mag:
                        sigma_.vector()[i] = 0
                    else:
                        sigma_.vector()[i] = v_mag*p_mag - alfa*p_mag**2 
        #----------------------------------------------------------------------------------------------------------------------------------
          
        # solve SUPG problem to get u_ 
        tau = Function(DGV)
        tau.vector()[:] = get_tau(v,eps)
        
        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                          ) # a way to realize swith on |v| == 0.
                        
        a = eps*inner(nabla_grad(u),nabla_grad(varphi))*dx + inner(v,nabla_grad(u))*varphi*dx +\
            tau*pert*inner(dot(nabla_grad(u),outer(v,v)),nabla_grad(varphi))*dx 
        
        l = f*varphi*dx + tau*pert*f*inner(v,nabla_grad(varphi))*dx

        A,L = assemble_system(a,l,bcs)
        solve(A,u_.vector(),L)

        # solve the non-linear problem                
        p_ = Function(DGVV)
        sigma_ = Function(DGV)
    
        aa = eps*inner(nabla_grad(u),nabla_grad(varphi))*dx + inner(v,nabla_grad(u))*varphi*dx+\
             tau*pert*inner(dot(nabla_grad(u),outer(v,v)),nabla_grad(varphi))*dx+\
             tau*sigma_*inner(nabla_grad(u),nabla_grad(varphi))*dx
            
        ll = f*varphi*dx + tau*pert*f*inner(v,nabla_grad(varphi))*dx
  
        error = 1
        tol = 1E-5
        iter = 0
        maxIter = 25
        
        u = Function(V)
        while error > tol and iter < maxIter:
            iter += 1
            
            projection(f,v,u_,p_)
            get_sigma(f,v,u_,p_,eps,sigma_)
            
            solve(aa==ll, u, bcs)
            
            diff = u.vector().array() - u_.vector().array()
            error = numpy.linalg.norm(diff, ord=numpy.Inf)
            u_.vector()[:] = u_.vector().array() + 1.0*diff        
            
            print iter, error
             
        plot(u_,interactive=True) 
        self.save(u,problem)
